#include <ros/ros.h>
#include <pcl_ros/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/registration/icp.h>
#include <pcl/registration/transforms.h>
#include <pcl/registration/registration.h>
#include <pcl/filters/passthrough.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/features/normal_3d.h>

#include <vector>
#include <iostream>
#include <algorithm>

#define POINTLIMIT 1e10
#define COARSE_LEAF_SIZE 0.05
#define FINE_LEAF_SIZE 0.001

//ICP Parameters
#define MAX_ICP_ITERATIONS 500
#define MAX_ICP_ERROR 1e-3
#define MIN_ICP_CORRESPONDENCES 1
#define ICP_OVERLAP_RATIO 0.8

#define DOWNSAMPLEREGISTRATION 1

using namespace std;

typedef pcl::PointXYZRGB PointType;
typedef pcl::PointCloud<PointType> kinectPCL;
typedef pcl::PointCloud<pcl::Normal> NormalCloud;

kinectPCL::Ptr fullModel(new kinectPCL);
kinectPCL::Ptr lastScan(new kinectPCL);
int nScans = 0;
//Use the last transform as the best guess for the next transform
Eigen::Matrix4f lastTransform = Eigen::Matrix4f::Identity();

ros::Publisher* debugPub;//Used to publish intermediate ICP scans
ros::Publisher* fullModelPub;//Used to publish the full model

void printPoints(kinectPCL::Ptr& cloud, int N) {
	kinectPCL::iterator iter = cloud->begin();
	int num = 0;
	while (iter != cloud->end()) {
		if (N > 0 && num >= N)
			break;
		PointType P = *iter;
		printf("(%g, %g, %g)\n", P.x, P.y, P.z);
		iter++;
		num++;
	}
}

int getNPoints(kinectPCL::Ptr& cloud) {
	return (int)(cloud->end() - cloud->begin());
}

void DownSample(kinectPCL::Ptr& cloud, kinectPCL::Ptr& downsampled, double res) {
	pcl::VoxelGrid<PointType> grid;
	grid.setInputCloud(cloud);
	grid.setLeafSize(res, res, res);
	grid.filter(*downsampled);
}

double getRMSD(kinectPCL::Ptr& cloud1, kinectPCL::Ptr& cloud2, pcl::KdTreeFLANN<PointType>& searchTree2) {
	double RMSD = 0.0;
	vector<int> NNIndex(1);
	vector<float> NNDist(1);
	kinectPCL::iterator iter = cloud1->begin();
	while (iter != cloud1->end()) {
		PointType P1 = *iter;
		searchTree2.nearestKSearch(P1, 1, NNIndex, NNDist);
		PointType P2 = cloud2->points[NNIndex[0]];
		double dx = P1.x - P2.x;
		double dy = P1.y - P2.y;
		double dz = P1.z - P2.z;
		RMSD += dx*dx + dy*dy + dz*dz;
	}
	double N = (double)getNPoints(cloud1);
	return sqrt(RMSD / N);
}

double getRMSD(kinectPCL::Ptr& modelCorrespPoints, kinectPCL::Ptr& targetCorrespPoints, int NPoints) {
	double RMSD = 0.0;
	for (int i = 0; i < NPoints; i++) {
		PointType P1 = modelCorrespPoints->points[i];
		PointType P2 = targetCorrespPoints->points[i];
		double dx = P1.x - P2.x;
		double dy = P1.y - P2.y;
		double dz = P1.z - P2.z;
		RMSD += dx*dx + dy*dy + dz*dz;
	}
	return sqrt(RMSD / (double)NPoints);
}

//Project P onto the plane defined by P0 and normal N
PointType projectPoint(PointType& P, PointType& P0, pcl::Normal& N) {
	double Nx = N.normal[0], Ny = N.normal[1], Nz = N.normal[2];
	if (Nx != Nx) //Normal is undefined for some reason
		return P0;
	double dx = P.x - P0.x, dy = P.y - P0.y, dz = P.z - P0.z;
	double NMagSqr = Nx*Nx + Ny*Ny + Nz*Nz;
	double projMag = (Nx*dx + Ny*dy + Nz*dz) / NMagSqr;
	//Subtract of component in normal direction
	dx -= projMag*Nx;
	dy -= projMag*Ny;
	dz -= projMag*Nz;
	PointType ret = P0;
	ret.x += dx;
	ret.y += dy;
	ret.z += dz;
	return ret;
}

class PointCorresp {
public:
	int modelIndex;
	int targetIndex;
	float dist;
	bool operator() (const PointCorresp& p1, const PointCorresp& p2) {
		return (p1.dist < p2.dist);
	}
};

//Align model to target using ICP with intial guess "initialGuess"
//and store the transformation that took the point clouds from its original
//frame to the final frame in "transform"
//Return whether or not ICP converged
bool ICPAlign(kinectPCL::Ptr& model, kinectPCL::Ptr& target, Eigen::Matrix4f& initialGuess, Eigen::Matrix4f* transform) {
	*transform = initialGuess;
	transformPointCloud(*model, *model, *transform);

	int Nmodel = getNPoints(model);
	int NSubset = (int)(ICP_OVERLAP_RATIO * (double)Nmodel);
	
	
	pcl::NormalEstimation<PointType, pcl::Normal> ne;
	ne.setInputCloud(target);
	pcl::KdTreeFLANN<PointType>::Ptr normalSearchTree(new pcl::KdTreeFLANN<PointType>());
	ne.setSearchMethod(normalSearchTree);
	NormalCloud::Ptr target_normals(new NormalCloud);
	ne.setRadiusSearch(COARSE_LEAF_SIZE*3);
	ne.compute(*target_normals);
	
	//Set up the KD tree used to search the target for nearest neightbors
	//and allocate space for the correspondence indices
	pcl::KdTreeFLANN<PointType> searchTree;
	searchTree.setInputCloud(target);	
	vector<int> NNIndex(1);
	vector<float> NNDist(1);
	kinectPCL::Ptr modelCorrespPoints(new kinectPCL);
	kinectPCL::Ptr targetCorrespPoints(new kinectPCL);
	if (NSubset < 1)
		return false;
	modelCorrespPoints->points.resize(NSubset);
	targetCorrespPoints->points.resize(NSubset);
	PointCorresp* initialCorresp = new PointCorresp[Nmodel];
	PointCorresp* lastCorresp = new PointCorresp[Nmodel];
	
	//ICP Iterations
	int iterations = 0;
	bool converged = false;
	double RMSD = FLT_MAX;
	while (!converged && iterations < MAX_ICP_ITERATIONS) {
		//Step 1: Find the point correspondences
		kinectPCL::iterator iter = model->begin();
		int i = 0;
		while (iter != model->end()) {
			PointType P = *iter;
			if (searchTree.nearestKSearch(P, 1, NNIndex, NNDist) > 0) {
				int targetIndex = NNIndex[0];
				float dist = NNDist[0];
				initialCorresp[i].modelIndex = i;
				initialCorresp[i].targetIndex = targetIndex;
				initialCorresp[i].dist = dist;
				/*PointType P2 = target->points[targetIndex];
				float dx = P2.x - P.x;
				float dy = P2.y - P.y;
				float dz = P2.z - P.z;
				float mydist = dx*dx + dy*dy + dz*dz;
				cout << dist << ": " << mydist << endl;*/
			}	
			iter++;
			i++;
		}
		//Step 2: Sort the correspondence in ascending order of distance
		//and only keep the closest half
		sort(initialCorresp, initialCorresp + Nmodel, initialCorresp[0]);
		for (int i = 0; i < NSubset; i++) {
			int modelIndex = initialCorresp[i].modelIndex;
			int targetIndex = initialCorresp[i].targetIndex;
			modelCorrespPoints->points[i] = model->points[modelIndex];
			//targetCorrespPoints->points[i] = target->points[targetIndex];
			//Project the model point onto the plane
			pcl::Normal normal = target_normals->points[targetIndex];
			targetCorrespPoints->points[i] = projectPoint(model->points[modelIndex], target->points[targetIndex], normal);
			PointType P0 = target->points[targetIndex];
			PointType P1 = targetCorrespPoints->points[i];
		}
		
		//Step 3: Find the optimal transformation between the points
		Eigen::Matrix4f nextTrans;
		estimateRigidTransformationSVD(*modelCorrespPoints, *targetCorrespPoints, nextTrans);
		transformPointCloud(*model, *model, nextTrans);
		debugPub->publish(model);//Publish the intermediate transformations for debugging
		(*transform) = (*transform)*nextTrans;
		
		//Step 4: Check for convergence
		//(the point correspondences haven't changed since the last time)
		bool sameCorresps = true;
		for (int k = 0; k < NSubset; k++) {
			if (lastCorresp[k].modelIndex != initialCorresp[k].modelIndex ||
				lastCorresp[k].targetIndex != initialCorresp[k].targetIndex) {
				sameCorresps = false;
				break;
			}
		}
		if (sameCorresps) {
			converged = true;
			break;
		}
		
		//Step 5: Save the correspondences from this time to check for convergence next time
		for (int k = 0; k < NSubset; k++) {
			lastCorresp[k] = initialCorresp[k];
		}
		iterations++;
	}
	
	RMSD = getRMSD(modelCorrespPoints, targetCorrespPoints, NSubset);
	
	cout << "\nICP " << (converged?"converged":"didn't converge") << " after " << iterations << " iterations with RMSD " << RMSD  << "\n\n";
	
	//Free up correspondence arrays
	delete[] initialCorresp;
	delete[] lastCorresp;
	
	return converged;
}

void PointsCallback(const kinectPCL::ConstPtr& scan) {
	kinectPCL::Ptr nextCloud(new kinectPCL);
	pcl::PassThrough<PointType> NaNFilter;
	NaNFilter.setInputCloud(scan);
	NaNFilter.setFilterFieldName("x");
	NaNFilter.setFilterLimits(-POINTLIMIT, POINTLIMIT);
	NaNFilter.filter(*nextCloud);
	if (fullModel->size() == 0) {
		*fullModel = *nextCloud;
		*lastScan = *nextCloud;
	}
	else {
		Eigen::Matrix4f transform;
		if (DOWNSAMPLEREGISTRATION) {
			//Do ICP on a downsampled version of the scans
			kinectPCL::Ptr toAlign(new kinectPCL);
			kinectPCL::Ptr alignTarget(new kinectPCL);
			DownSample(nextCloud, toAlign, COARSE_LEAF_SIZE);
			//Align pairwise scan
			//DownSample(fullModel, alignTarget, COARSE_LEAF_SIZE);
			DownSample(lastScan, alignTarget, COARSE_LEAF_SIZE);			
			
			printf("Full Model #Points: %i\n", getNPoints(fullModel));
			printf("Next Cloud #Points: %i\n", getNPoints(nextCloud));
			printf("Downsampled #Points: %i\n", getNPoints(toAlign));
	
			//Do ICP with the last transformation as the intial guess
			ICPAlign(toAlign, alignTarget, lastTransform, &transform);
			//Now snap the full version of the point clouds in place
			transformPointCloud(*nextCloud, *nextCloud, transform);
		}
		else {
			//Do ICP on the full model
			//ICPAlign(nextCloud, fullModel, lastTransform, &transform);
			ICPAlign(nextCloud, lastScan, lastTransform, &transform);
		}		
		*fullModel += *nextCloud;
		/*if (nScans%10 == 0) {
			//Every tenth scan downsample the full model
			kinectPCL::Ptr nextFullModel(new kinectPCL);
			DownSample(fullModel, nextFullModel, FINE_LEAF_SIZE);
			*fullModel = *nextFullModel;
		}*/
		//Update the working transform
		lastTransform = transform;
		*lastScan = *nextCloud;
	}
	fullModelPub->publish(fullModel);
	nScans++;
}

int main(int argc, char** argv) {
    ros::init(argc, argv, "testconversion");
    ros::NodeHandle n;
    
	if (argc < 3) {
		fprintf(stderr, "ERROR: Did not specify sourceCloud and fullModelCloud in \"args\"");
		return 1;
	}
	char* sourceCloud = argv[1];
	char* fullModelCloud = argv[2];
    
    ros::Subscriber rfidSub = n.subscribe<kinectPCL>(sourceCloud, 100, PointsCallback);
    ros::Publisher pub = n.advertise<kinectPCL>(fullModelCloud, 20);
    ros::Publisher pub2 = n.advertise<kinectPCL>("tralierfpc/debug", 2);
    fullModelPub = &pub;
    debugPub = &pub2;
    ros::Rate loop_rate(5);

    while(n.ok()){
        ros::spinOnce();
        pub.publish(fullModel);
        loop_rate.sleep();
    }
    return 0;
}
